import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import matplotlib as mpl

mpl.use('TkAgg')

data = pd.read_csv("data/Q3A_pred.csv")
del data['date']
print(data)

df = pd.read_csv("data/Q2.csv", index_col=0)
# 将df序号设置为date
# df.index = df['date']
del df['date']
df = df.iloc[:, -1:]
n_train = round(df.shape[0] * 0.8)
train = df.iloc[:n_train, :]    # 2433
test = df.iloc[n_train:, :]     # 608 4*152  16*38
test = test.reset_index()
del test['index']

data.index.name = "Days"
test.index.name = "Days"

data_list = data["AQI"].values
test_list = test["AQI"].values

print(data_list)
print(test_list)


# 绘制预测/真实值图

# plt.figure(figsize=(8, 6), dpi=300)
#
# plt.plot(test.index,  test["AQI"], label='observed')
# plt.plot(data.index, data["AQI"], label='forecast')
#
# plt.legend()
# plt.savefig("img/Q3_pic1.png")
# plt.show()

def NMSE(list1, list2):
    list1 = np.array(list1)
    list2 = np.array(list2)
    # 计算均方误差MSE
    mse = np.mean((list1 - list2) ** 2)
    # 计算根均方误差RMSE
    rmse = np.sqrt(mse)
    return rmse

ans1 = NMSE(data_list, test_list)
print(ans1)

data_preall = pd.read_csv("data/Q3A_pred1.csv")
data_preall.index = data_preall['date']
del data_preall['date']
print(data_preall)
data_all = pd.read_csv("data/Q2.csv", index_col=0)
data_all.index = data_all['date']
del data_all['date']
data_all = data_all.iloc[:, -1:]
print(data_all)


data_preall_list = data_preall["pred"].values
data_all_list = data_all["AQI"].values
print(NMSE(data_all_list, data_preall_list))

# 仅用ARIMA大概是31左右